import pandas as pd
import jieba
import nltk
import wordcloud
import jieba.posseg as psg
import imageio
import matplotlib.pyplot as plt
from wordcloud import get_single_color_func
from wordcloud import ImageColorGenerator
import logging
import numpy as np
jieba.setLogLevel(logging.INFO)


class GroupedColorFunc(object):

    def __init__(self, color_to_words, default_color):
        self.color_func_to_words = [
            (get_single_color_func(color), set(words))
            for (color, words) in color_to_words.items()]

        self.default_color_func = get_single_color_func(default_color)

    def get_color_func(self, word):
        """Returns a single_color_func associated with the word"""
        try:
            color_func = next(
                color_func for (color_func, words) in self.color_func_to_words
                if word in words)
        except StopIteration:
            color_func = self.default_color_func

        return color_func

    def __call__(self, word, **kwargs):
        return self.get_color_func(word)(word, **kwargs)


raw = pd.read_csv(r"D:\代码\nlp\金庸-射雕英雄传txt精校版.txt",
                  names = ['txt'], sep ='aaa', encoding ="utf-8" ,engine='python')


# 章节判断用变量预处理
def m_head(tmpstr):
    return tmpstr[:1]


def m_mid(tmpstr):
    return tmpstr.find("回 ")
raw['head'] = raw.txt.apply(m_head)
raw['mid'] = raw.txt.apply(m_mid)
raw['len'] = raw.txt.apply(len)

# 章节判断
chapnum = 0
for i in range(len(raw)):
    if raw['head'][i] == "第" and raw['mid'][i] > 0 and raw['len'][i] < 30:
        chapnum += 1
    if chapnum >= 40 and raw['txt'][i] == "附录一：成吉思汗家族":
        chapnum = 0
    raw.loc[i, 'chap'] = chapnum

# 删除临时变量
del raw['head']
del raw['mid']
del raw['len']

tmpchap = raw[raw['chap'] == 1].copy()
tmpchap.reset_index(drop=True, inplace=True)
tmpchap['paraidx'] = tmpchap.index


# 进行所求结果的输出
result_list = []
for i in range(len(tmpchap.index)):
    tmppara = tmpchap[tmpchap['paraidx'] == i].copy()
    tmpstr = tmppara.txt[i]
    result_list.append(tmpstr)


savelist1 = []   #地名
savelist2 = []   #人名
for list in result_list:
    for sentence in list:
        tmpres = psg.cut(sentence)  # 附加词性的分词结果
        # print(tmpres)
        for item in tmpres:
            # print(item.word, item.flag)
            if item.flag == 'ns':
                savelist1.append(item.word)
            elif item.flag == 'nr':
                savelist2.append(item.word)
savelist = savelist1+savelist2
print(savelist1)
print(savelist2)
freq = {}

for word in savelist:
    if word in freq:
        freq[word] += 1
    else:
        freq[word] = 1
# print(freq)

df = pd.DataFrame(savelist, columns = ['word'])
result = df.groupby(['word']).size()

freqlist = result.sort_values(ascending=False)

fdist = nltk.FreqDist(savelist) # 生成完整的词条频数字典




myfont = 'STLITI.TTF'
cloudobj = wordcloud.WordCloud(font_path = myfont,
    mask = imageio.imread("射雕背景.jpg"),
    mode = "RGBA", background_color = None
    ).generate(' '.join(savelist))
color_to_words = {
    'blue': savelist2,
    'red': savelist1
}


default_color = 'grey' # 指定其他词条的颜色

grouped_color_func = GroupedColorFunc(color_to_words, default_color)

cloudobj.recolor(color_func=grouped_color_func)
imgobj = imageio.imread('blue.jpg')
image_colors = wordcloud.ImageColorGenerator(np.array(imgobj))
cloudobj.recolor(color_func=image_colors)
cloudobj.to_file("蓝色词云图.png")
plt.imshow(cloudobj)
plt.axis("off")
plt.show()